#include "6-balance_bit_tree.hpp"
#include "7-red_black_node.hpp"

template<class T = int>
void printInOrderRbNode(RbNode<T>* p_node)
{
  if (p_node != nullptr) {
    int dep = getDepthTr(p_node);
    for (int i = 0; i < dep; ++i) {
      std::cout << "      " ;
    }
    std::cout << p_node->color << "-" << p_node->data << std::endl;
  }
}

template<class T = int>
void printRbTree(RbNode<T>* tree)
{
  if (tree == nullptr) {
    return;
  }

  printRbTree(dynamic_cast<RbNode<T>*>(tree->rchild));
  printInOrderRbNode(tree);
  printRbTree(dynamic_cast<RbNode<T>*>(tree->lchild));
}

// 在根为root的树中吸收红色节点target
template<class T = int>
BitNode<T>* absorbRed(BitNode<T>* root, BitNode<T>* target)
{
  if (root == nullptr || target == nullptr) {
    return root;
  }

  /* 节点为根节点，颜色变黑即可 */
  auto father = target->parent;
  if (father == nullptr) {
    std::cout << "--target-" << target->data << " is root--" << std::endl;
    setColor(target, BLACK);
    return root;
  }

  /* 父节点为黑色节点，对应2节点和特殊情况下的3节点，无需调整 */
  if (getColor(father) == BLACK) {
    std::cout << "--father-" << father->data << " is black node---" << std::endl;
    return root;
  }

  /* 父节点为红色节点 */
  auto grandpa = father->parent;
  auto far_is_lchild = (grandpa->lchild == father);
  auto uncle = far_is_lchild ? grandpa->rchild : grandpa->lchild;

  if (getColor(uncle) == BLACK) {
    // 父节点与祖父节点组成3 节点
    std::cout << "--father-" << father->data << " is 3 node---" << std::endl;

    BitNode<T>* &father_mount = far_is_lchild ? grandpa->lchild : grandpa->rchild;
    auto great_grandpa = grandpa->parent;
    BitNode<T>* &grandpa_mount = (great_grandpa == nullptr) ? root :
      ((great_grandpa->lchild == grandpa) ? great_grandpa->lchild : great_grandpa->rchild);

    setColor(grandpa, RED);
    if (far_is_lchild) {
      if (target == father->rchild) {
        setColor(target, BLACK);
        father_mount = leftRotate(father);  // LR
      } else {
        setColor(father, BLACK);
      }
      grandpa_mount = rightRotate(grandpa);  // LL
    } else {
      if (target == father->lchild) {
        setColor(target, BLACK);
        father_mount = rightRotate(father); // RL
      } else {
        setColor(father, BLACK);
      }
      grandpa_mount = leftRotate(grandpa);  // RR
    }
  } else {
    // 叔叔、父及祖父节点组成4 节点
    std::cout << "---father-" << father->data << " is 4 node---" << std::endl;

    setColor(grandpa, RED);
    setColor(father, BLACK);
    setColor(uncle, BLACK);
    return absorbRed(root, grandpa);
  }

  return root;
}

template<class T = int>
BitNode<T>* insertNodeRb(BitNode<T>* tree, T data)
{
  if (tree == nullptr) {
    auto target = new RbNode<T>(data);
    target->color = BLACK;
    return target;
  }

  // 查找节点的插入位置
  BitNode<T>* father = nullptr;
  BitNode<T>* target = tree;
  while (target != nullptr) {
    if (data < target->data) {
      father = target;
      target = target->lchild;
    } else if (data > target->data) {
      father = target;
      target = target->rchild;
    } else {
      break;
    }
  }
  // 不为空则节点已存在
  if (target != nullptr) {
    return tree;
  }

  target = new RbNode<T>(data);
  if (data < father->data) father->lchild = target;
  if (data > father->data) father->rchild = target;
  target->parent = father;

  return absorbRed(tree, target);
}

template<class T = int>
BitNode<T>* replenishBlack(BitNode<T>* root, BitNode<T>* father, bool tar_is_lchild)
{
  if (root == nullptr || father == nullptr) {
    std::cout << "--find to root--" << std::endl;
    return root;
  }

  // 找补黑色，兄弟节点必非空
  auto sib = tar_is_lchild ? father->rchild : father->lchild;
  if (getColor(sib) == RED) {
    // 此时的sib并非目标节点对应2-3-4树中的兄弟节点
    std::cout << "------father-" << father->data << "- is 3 black node---" << std::endl;

    auto grandpa = father->parent;
    BitNode<T>* &father_mount = (grandpa == nullptr) ? root :
      ((grandpa->lchild == father) ? grandpa->lchild : grandpa->rchild);

    setColor(sib, BLACK);
    setColor(father, RED);
    father_mount = tar_is_lchild ? leftRotate(father) : rightRotate(father);

    // 获取旋转后真正的兄弟节点
    sib = tar_is_lchild ? father->rchild : father->lchild;
  }

  auto grandpa = father->parent;
  BitNode<T>* &father_mount = (grandpa == nullptr) ? root :
    ((grandpa->lchild == father) ? grandpa->lchild : grandpa->rchild);

  auto father_color = getColor(father);
  auto sib_weight = nodeWeight(dynamic_cast<RbNode<T>*>(sib));

  if (sib_weight == 4) {
    // 兄弟节点为4 节点，通过旋转补充黑色
    std::cout << "----sib-" << sib->data << "- is 4 node---" << std::endl;

    setColor(father, BLACK);
    setColor(sib, father_color);
    if (tar_is_lchild) {
      setColor(sib->rchild, BLACK);
      father_mount = leftRotate(father);
    } else {
      setColor(sib->lchild, BLACK);
      father_mount = rightRotate(father);
    }
  } else if (sib_weight == 3) {
    // 兄弟节点为3 节点，通过旋转补充黑色
    std::cout << "----sib-" << sib->data << "- is 3 node---" << std::endl;

    setColor(father, BLACK);
    setColor(sib, father_color);
    if (tar_is_lchild) {
      if (getColor(sib->lchild) == RED) {
        setColor(sib->lchild, BLACK);
        father->rchild = rightRotate(sib);  // RL
      } else {
        setColor(sib->rchild, BLACK);
      }
      father_mount = leftRotate(father);  // RR
    } else {
      if (getColor(sib->rchild) == RED) {
        setColor(sib->rchild, BLACK);
        father->lchild = leftRotate(sib);  // LR
      } else {
        setColor(sib->lchild, BLACK);
      }
      father_mount = rightRotate(father);  // LL
    }
  } else {
    // 兄弟节点为2 节点，染红后将父节点染黑
    std::cout << "----sib-" << sib->data << "- is 2 node---" << std::endl;

    setColor(sib, RED);
    if (father_color == BLACK) {
      // 父节点为黑色，递归补充黑色
      std::cout << "------father-" << father->data << "- is black node---" << std::endl;

      // 获取父节点属于哪个分支
      auto grandpa = father->parent;
      auto far_is_lchild = (grandpa == nullptr) ? false : (grandpa->lchild == father);
      return replenishBlack(root, grandpa, far_is_lchild);
    }

    std::cout << "------father-" << father->data << "- is red node---" << std::endl;
    setColor(father, BLACK);
  }
  return root;
}

template<class T = int>
BitNode<T>* delNodeRb(BitNode<T>* root, BitNode<T>* target)
{
  if (root == nullptr || target == nullptr) {
    return root;
  }
  /* 调整后的target只能为2、3节点及4节点的副节点 */

  auto father = target->parent;

  auto tar_is_lchild = (father == nullptr) ? false : (father->lchild == target);
  BitNode<T>* &mount = (father == nullptr) ? root : (tar_is_lchild ? father->lchild : father->rchild);

  /* 节点为红，则为3、4节点的副节点 */
  if (getColor(target) == RED) {
    std::cout << "--target-" << target->data << " is red node---" << std::endl;
    free(target);
    mount = nullptr;
    return root;
  }

  /* 节点为黑、有孩子，则为3 节点的上黑节点 */
  auto child = (target->lchild != nullptr) ? target->lchild : target->rchild;
  if (child != nullptr) {
    std::cout << "--target-" << target->data << " is 3 black node---" << std::endl;
    setParent(child, father);
    setColor(child, BLACK);

    free(target);
    mount = child;
    return root;
  }

  /* 节点为黑、无孩子，则为2 节点 */
  std::cout << "--target-" << target->data << " is 2 node---" << std::endl;
  free(target);
  mount = nullptr;

  return replenishBlack(root, father, tar_is_lchild);
}

template<class T = int>
BitNode<T>* delNodeRb(BitNode<T>* tree, T data)
{
  auto node = searchNodeSTr(tree, data);
  auto target = searchPrecursor(node);
  //auto target = searchSuccessor(node);
  if (node == nullptr || target == nullptr) {
    std::cout << "-target not find---" << std::endl;
    return tree;
  }
  
  T temp = node->data;
  node->data = target->data;
  target->data = temp;

  /* 节点交换后，待删除节点为2-3-4树的叶子节点 */
  return delNodeRb(tree, target);
}

template<class T = int>
class RbTree : public BalanceBitTree<T> {
  public:
    RbTree() = default;
    explicit RbTree(RbNode<T>* _tree) : BalanceBitTree<T>(_tree) {}
    virtual ~RbTree() = default;

    RbNode<T>* operator->() { return dynamic_cast<RbNode<T>*>(BalanceBitTree<T>::tree); }

    int getHight() override { return getHightRb(dynamic_cast<RbNode<T>*>(BalanceBitTree<T>::tree)); }

    void insertNode(T data) override { BalanceBitTree<T>::tree = insertNodeRb(BalanceBitTree<T>::tree, data); }
    void delNode(T data) override { BalanceBitTree<T>::tree = delNodeRb(BalanceBitTree<T>::tree, data); }

    void print() override { 
      printRbTree(dynamic_cast<RbNode<T>*>(BalanceBitTree<T>::tree));
      std::cout << "============================================================" << std::endl;
    }
};

